'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: yolov3.py
@time: 2020-06-19 11:04:38
@desc: 
'''
import torch
from jjzhk.config import ZKCFG
from jjzhk.device import device
import numpy as np
from ELib.utils.utils_yolov3 import to_cpu, non_max_suppression, build_targets, rescale_boxes
from .model_zoo import MODEL_ZOO
from .segbasemodel import SegBaseModel


@MODEL_ZOO.register()
class Darknet(SegBaseModel):
    def __init__(self,cfg:ZKCFG):
        super().__init__(cfg, need_backbone=False)
        layers = self._create_layers_()
        self.model_def_layer, self.model_def = self.create_model(layers)
        self.module_list = torch.nn.ModuleList(self.model_def_layer)
        # print(self.module_list)

        self.seen = 0
        self.header_info = np.array([0, 0, 0, self.seen, 0], dtype=np.int32)

    def _create_layers_(self):
        anchor_idxs = [int(x) for x in self.cfg.MODEL.YOLOLAYER_01.MASK.split(",")]
        anchors = [int(x) for x in self.cfg.MODEL.YOLOLAYER_01.ANCHORS.split(",")]
        anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
        yololayer01_anchors = [anchors[i] for i in anchor_idxs]

        anchor_idxs = [int(x) for x in self.cfg.MODEL.YOLOLAYER_02.MASK.split(",")]
        anchors = [int(x) for x in self.cfg.MODEL.YOLOLAYER_02.ANCHORS.split(",")]
        anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
        yololayer02_anchors = [anchors[i] for i in anchor_idxs]

        anchor_idxs = [int(x) for x in self.cfg.MODEL.YOLOLAYER_03.MASK.split(",")]
        anchors = [int(x) for x in self.cfg.MODEL.YOLOLAYER_03.ANCHORS.split(",")]
        anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
        yololayer03_anchors = [anchors[i] for i in anchor_idxs]

        layers = [
            [
                torch.nn.Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(32, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(32, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1,1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(64, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],
            [
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(1024, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(1024, 3 * (self.cfg.BASE.NUM_CLASSES + 5), kernel_size=(1, 1), stride=(1, 1))
            ],[
                YOLOLayer(yololayer01_anchors, self.cfg.BASE.NUM_CLASSES, self.cfg.BASE.IMAGE_SIZE[1])
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                Upsample(scale_factor=2)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(512, 3 * (self.cfg.BASE.NUM_CLASSES + 5), kernel_size=(1, 1), stride=(1, 1))
            ],[
                YOLOLayer(yololayer02_anchors, self.cfg.BASE.NUM_CLASSES, self.cfg.BASE.IMAGE_SIZE[1])
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                Upsample(scale_factor=2)
            ],[
                EmptyLayer()
            ],[
                torch.nn.Conv2d(384, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],
            [
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False),
                torch.nn.BatchNorm2d(128, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
                torch.nn.BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True),
                torch.nn.LeakyReLU(negative_slope=0.1)
            ],[
                torch.nn.Conv2d(256, 3 * (self.cfg.BASE.NUM_CLASSES + 5), kernel_size=(1, 1), stride=(1, 1))
            ],[
                YOLOLayer(yololayer03_anchors, self.cfg.BASE.NUM_CLASSES, self.cfg.BASE.IMAGE_SIZE[1])
            ]
        ]

        return layers

    def create_model(self, layers):
        modeldef_layers = []
        model_def = []
        for index, layer_model in enumerate(layers):
            seq = torch.nn.Sequential()
            def_name = []
            for layer in layer_model:
                name = layer.__class__.__name__
                seq_layer_name = ""
                if name == 'Conv2d':
                    seq_layer_name = 'conv'
                    def_name.append('conv')
                elif name == 'BatchNorm2d':
                    seq_layer_name = 'batch_norm'
                    def_name.append('norm')
                elif name == 'LeakyReLU':
                    seq_layer_name = 'leaky'
                    def_name.append('leaky')
                elif name == 'EmptyLayer':
                    if index == 83 or index == 86 or index == 95 or index == 98:
                        seq_layer_name = 'route'
                        def_name.append('route')
                    else:
                        seq_layer_name = 'shortcut'
                        def_name.append('shortcut')
                elif name == 'YOLOLayer':
                    seq_layer_name = 'yolo'
                    def_name.append('yolo')
                elif name == 'Upsample':
                    seq_layer_name = 'upsample'
                    def_name.append('upsample')

                seq.add_module("%s_%d" % (seq_layer_name, index), layer)

            model_def.append(def_name)
            modeldef_layers += [seq]

        return modeldef_layers, model_def

    def forward(self, x, target=None, **kwargs):
        img_dim = x.shape[2]
        loss = 0
        route = ["-4", "-1, 61", "-4", "-1, 36"]
        route_index = 0
        layer_outputs, yolo_outputs = [], []
        for i, (module, module_def) in enumerate(zip(self.module_list, self.model_def)):
            if 'conv' in module_def or "upsample" in module_def or "maxpool" in module_def:
                x = module(x)
            elif "route" in module_def:
                x = torch.cat([layer_outputs[int(layer_i)] for layer_i in route[route_index].split(",")], 1)
                route_index += 1
            elif "shortcut" in module_def:
                layer_i = int(-3)
                x = layer_outputs[-1] + layer_outputs[layer_i]
            elif "yolo" in module_def:
                x, layer_loss = module[0](x, target, img_dim)
                loss += layer_loss / x.shape[0]
                yolo_outputs.append(x)
            layer_outputs.append(x)
        yolo_outputs = to_cpu(torch.cat(yolo_outputs, 1))
        return yolo_outputs if target is None else loss

    def get_detections(self,image, **kwargs):
        image = torch.from_numpy(image).float()

        with torch.no_grad():
            image = torch.autograd.Variable(image)

        image = image.to(device)  # 1, 10647, 85
        detections = self(image)
        detections = non_max_suppression(detections, self.cfg.BASE.CONF_THRESHOLD, self.cfg.BASE.IOU_THRESHOLD)

        return detections

    def get_eval_predictions(self,info, detections):
        re_boxes = [[] for _ in range(len(self.cfg.BASE.CLASSINFO.keys()) + 1)]

        width = info[0]['width']
        height = info[0]['height']
        detections = rescale_boxes(detections[0], self.cfg.BASE.IMAGE_SIZE[0], (height, width)).cpu().numpy()

        for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
            cls_dets = np.array([x1, y1, x2, y2, cls_conf], dtype=np.float32)
            re_boxes[int[conf]] = cls_dets

        return re_boxes

    def get_predict(self, image, info, **kwargs):
        detections = self.get_detections(image, **kwargs)

        width = info[0]['width']
        height = info[0]['height']
        class_names = self.cfg.BASE.CLASSINFO.keys()
        detections = rescale_boxes(detections[0], self.cfg.BASE.IMAGE_SIZE[0], (height, width)).cpu().numpy()

        reBoxes = []

        for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
            reBoxes.append([(x1, y1), (x2, y2), class_names[int(cls_pred)], "", cls_conf])
        return [reBoxes]

    # def load_init_weights(self, weights_path):
    #     """Parses and loads the weights stored in 'weights_path'"""
    #
    #     # Open the weights file
    #     with open(weights_path, "rb") as f:
    #         header = np.fromfile(f, dtype=np.int32, count=5)  # First five are header values
    #         self.header_info = header  # Needed to write header when saving weights
    #         self.seen = header[3]  # number of images seen during training
    #         weights = np.fromfile(f, dtype=np.float32)  # The rest are weights
    #
    #     # Establish cutoff for loading backbone weights
    #     cutoff = None
    #     if "darknet53.conv.74" in weights_path:
    #         cutoff = 75
    #
    #     ptr = 0
    #     for i, (module_def, module) in enumerate(zip(self.model_def, self.module_list)):
    #         if i == cutoff:
    #             break
    #         if "conv" in module_def:
    #             conv_layer = module[0]
    #             if "norm" in module_def:
    #                 # Load BN bias, weights, running mean and running variance
    #                 bn_layer = module[1]
    #                 num_b = bn_layer.bias.numel()  # Number of biases
    #                 # Bias
    #                 bn_b = torch.from_numpy(weights[ptr: ptr + num_b]).view_as(bn_layer.bias)
    #                 bn_layer.bias.data.copy_(bn_b)
    #                 ptr += num_b
    #                 # Weight
    #                 bn_w = torch.from_numpy(weights[ptr: ptr + num_b]).view_as(bn_layer.weight)
    #                 bn_layer.weight.data.copy_(bn_w)
    #                 ptr += num_b
    #                 # Running Mean
    #                 bn_rm = torch.from_numpy(weights[ptr: ptr + num_b]).view_as(bn_layer.running_mean)
    #                 bn_layer.running_mean.data.copy_(bn_rm)
    #                 ptr += num_b
    #                 # Running Var
    #                 bn_rv = torch.from_numpy(weights[ptr: ptr + num_b]).view_as(bn_layer.running_var)
    #                 bn_layer.running_var.data.copy_(bn_rv)
    #                 ptr += num_b
    #             else:
    #                 # Load conv. bias
    #                 num_b = conv_layer.bias.numel()
    #                 conv_b = torch.from_numpy(weights[ptr: ptr + num_b]).view_as(conv_layer.bias)
    #                 conv_layer.bias.data.copy_(conv_b)
    #                 ptr += num_b
    #             # Load conv. weights
    #             num_w = conv_layer.weight.numel()
    #             conv_w = torch.from_numpy(weights[ptr: ptr + num_w]).view_as(conv_layer.weight)
    #             conv_layer.weight.data.copy_(conv_w)
    #             ptr += num_w


class EmptyLayer(torch.nn.Module):
    """Placeholder for 'route' and 'shortcut' layers"""
    def __init__(self):
        super(EmptyLayer, self).__init__()


class Upsample(torch.nn.Module):
    """ nn.Upsample is deprecated """

    def __init__(self, scale_factor, mode="nearest"):
        super(Upsample, self).__init__()
        self.scale_factor = scale_factor
        self.mode = mode

    def forward(self, x):
        x = torch.nn.functional.interpolate(x, scale_factor=self.scale_factor, mode=self.mode)
        return x


class YOLOLayer(torch.nn.Module):
    """Detection layer"""

    def __init__(self, anchors, num_classes, img_dim=416):
        super(YOLOLayer, self).__init__()
        self.anchors = anchors
        self.num_anchors = len(anchors)
        self.num_classes = num_classes
        self.ignore_thres = 0.5
        self.mse_loss = torch.nn.MSELoss()
        self.bce_loss = torch.nn.BCELoss()
        self.obj_scale = 1
        self.noobj_scale = 100
        self.metrics = {}
        self.img_dim = img_dim
        self.grid_size = 0  # grid size

    def compute_grid_offsets(self, grid_size, cuda=True):
        self.grid_size = grid_size
        g = self.grid_size
        FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
        self.stride = self.img_dim / self.grid_size
        # Calculate offsets for each grid
        self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor)
        self.grid_y = torch.arange(g).repeat(g, 1).t().view([1, 1, g, g]).type(FloatTensor)
        self.scaled_anchors = FloatTensor([(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors])
        self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))
        self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))

    def forward(self, x, targets=None, img_dim=None):
        '''
        targets的shape是(boxes number, 6)。boxes number是这batch size个images所包含的所有boxes的个数.
        6: 分别是center_x, center_y, width, height,Conf, cls尺寸是归一化只有的尺寸
        x的shape是不确定的，因为输入的图片的size不确定。对于image size=416的时候，有:
        # 1. x => batch_size, 78, 13, 13 target => boxes number, 6
        # 2. x => batch_size, 78, 26, 26 target => boxes number, 6
        # 3. x => batch_size, 78, 52, 52 target => boxes number, 6
        '''
        # Tensors for cuda support
        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor

        self.img_dim = img_dim
        num_samples = x.size(0)
        grid_size = x.size(2)

        prediction = (
            x.view(num_samples, self.num_anchors, self.num_classes + 5, grid_size, grid_size)
            .permute(0, 1, 3, 4, 2)
            .contiguous()
        )

        # Get outputs
        '''
        predictions是预测值，我们需要将center的x,y两个坐标归化到0~1之间，于是使用了sigmoid
        如果不使用sigmoid，那么在训练前期会很难收敛
        '''
        x = torch.sigmoid(prediction[..., 0])  # Center x \sigma{x}
        y = torch.sigmoid(prediction[..., 1])  # Center y \sigma{y}
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height
        pred_conf = torch.sigmoid(prediction[..., 4])  # Conf
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

        # If grid size does not match current we compute new offsets
        if grid_size != self.grid_size:
            self.compute_grid_offsets(grid_size, cuda=x.is_cuda)

        # Add offset and scale with anchors
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + self.grid_x # \sigma{t_x} + cx
        pred_boxes[..., 1] = y.data + self.grid_y # \sigma{t_y} + cy
        pred_boxes[..., 2] = torch.exp(w.data) * self.anchor_w # \alpha_w \times e^{t_w}
        pred_boxes[..., 3] = torch.exp(h.data) * self.anchor_h # \alpha_h \times e^{t_h}

        output = torch.cat(
            (
                pred_boxes.view(num_samples, -1, 4) * self.stride,
                pred_conf.view(num_samples, -1, 1),
                pred_cls.view(num_samples, -1, self.num_classes),
            ),
            -1,
        )

        if targets is None:
            return output, 0
        else:
            iou_scores, class_mask, obj_mask, noobj_mask, tx, ty, tw, th, tcls, tconf = build_targets(
                pred_boxes=pred_boxes,
                pred_cls=pred_cls,
                target=targets,
                anchors=self.scaled_anchors,
                ignore_thres=self.ignore_thres,
            )

            # Loss : Mask outputs to ignore non-existing objects (except with conf. loss)
            obj_mask = obj_mask.bool()

            loss_x = self.mse_loss(x[obj_mask], tx[obj_mask])
            loss_y = self.mse_loss(y[obj_mask], ty[obj_mask])
            loss_w = self.mse_loss(w[obj_mask], tw[obj_mask])
            loss_h = self.mse_loss(h[obj_mask], th[obj_mask])
            loss_conf_obj = self.bce_loss(pred_conf[obj_mask], tconf[obj_mask])

            noobj_mask = noobj_mask.bool()
            loss_conf_noobj = self.bce_loss(pred_conf[noobj_mask], tconf[noobj_mask])
            loss_conf = self.obj_scale * loss_conf_obj + self.noobj_scale * loss_conf_noobj
            loss_cls = self.bce_loss(pred_cls[obj_mask], tcls[obj_mask])
            total_loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls

            # Metrics
            cls_acc = 100 * class_mask[obj_mask].mean()
            conf_obj = pred_conf[obj_mask].mean()
            conf_noobj = pred_conf[noobj_mask].mean()
            conf50 = (pred_conf > 0.5).float()
            iou50 = (iou_scores > 0.5).float()
            iou75 = (iou_scores > 0.75).float()
            detected_mask = conf50 * class_mask * tconf
            precision = torch.sum(iou50 * detected_mask) / (conf50.sum() + 1e-16)
            recall50 = torch.sum(iou50 * detected_mask) / (obj_mask.sum() + 1e-16)
            recall75 = torch.sum(iou75 * detected_mask) / (obj_mask.sum() + 1e-16)

            self.metrics = {
                "loss": to_cpu(total_loss).item(),
                "x": to_cpu(loss_x).item(),
                "y": to_cpu(loss_y).item(),
                "w": to_cpu(loss_w).item(),
                "h": to_cpu(loss_h).item(),
                "conf": to_cpu(loss_conf).item(),
                "cls": to_cpu(loss_cls).item(),
                "cls_acc": to_cpu(cls_acc).item(),
                "recall50": to_cpu(recall50).item(),
                "recall75": to_cpu(recall75).item(),
                "precision": to_cpu(precision).item(),
                "conf_obj": to_cpu(conf_obj).item(),
                "conf_noobj": to_cpu(conf_noobj).item(),
                "grid_size": grid_size,
            }

            return output, total_loss
